package com.example.alibabanacosdiscoveryclient01.controller;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.postretrieval.compression.DocumentCompressor;
import org.springframework.ai.rag.postretrieval.ranking.DocumentRanker;
import org.springframework.ai.rag.postretrieval.selection.DocumentSelector;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.CompressionQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
import org.springframework.ai.rag.retrieval.join.ConcatenationDocumentJoiner;
import org.springframework.ai.rag.retrieval.join.DocumentJoiner;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author:guoq
 * @date:2025/2/11
 * @descripion:
 */
@RestController
@RequestMapping("/ai")
@Tag(name = "RAG高级应用")
public class RAGController {
    @Autowired
    VectorStore vectorStore;
    @Autowired
    ChatClient.Builder builder;

    private final ChatClient chatClient;

    public RAGController(ChatClient.Builder builder) {
        //配置AI助手角色
        this.chatClient = builder
                .defaultSystem("你是一位专业智能助手")
                .build();
    }

    @GetMapping("/rag/pre-cqt")
    @Operation(summary = "rag预处理-上下文感知查询")
    public String rag(@RequestParam(defaultValue = "有哪些区") String message) {
        Query query = Query.builder()
                .text(message)
                .history(new UserMessage("中国的首都是哪里"),
                        new AssistantMessage("中国的首都是北京"))
                .build();
        //创建查询转换器，将带有上下文的查询转换成完整的独立查询
        CompressionQueryTransformer queryTransformer = CompressionQueryTransformer.builder()
                .chatClientBuilder(builder)
                .build();
        //执行查询转换
        Query transform = queryTransformer.transform(query);
        String userMsg = transform.text();
        System.out.println(userMsg);
        //调用大模型
        String content = chatClient.prompt(new Prompt(userMsg))
                .call()
                .content();
        return content;
    }

    @GetMapping("/rag/pre-rqt")
    @Operation(summary = "rag预处理-查询重写")
    public String rag2(@RequestParam(defaultValue = "我正在学人工智能，什么是RAG?") String message) {
        Query query = new Query(message);
        RewriteQueryTransformer queryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(builder)
                .build();
        Query transformedQuery = queryTransformer.transform(query);
        String userMsg = transformedQuery.text();
        System.out.println(userMsg);
        //调用大模型
        String content = chatClient.prompt(new Prompt(userMsg))
                .call()
                .content();
        return content;
    }

    @GetMapping("/rag/pre-tqt")
    @Operation(summary = "rag预处理-查询翻译")
    public String rag3(@RequestParam(defaultValue = "what is RAG?") String message) {
        Query query = new Query(message);
        TranslationQueryTransformer queryTransformer = TranslationQueryTransformer.builder()
                .chatClientBuilder(builder)
                .targetLanguage("japanese")
                .build();
        Query transformedQuery = queryTransformer.transform(query);
        String userMsg = transformedQuery.text();
        System.out.println(userMsg);
        //调用大模型
        String content = chatClient.prompt(new Prompt(userMsg))
                .call().content();
        return content;
    }

    @GetMapping("/rag/pre-mqe")
    @Operation(summary = "rag预处理-查询扩展")
    public String rag4(@RequestParam(defaultValue = "什么是RAG?") String message) {
        MultiQueryExpander expander = MultiQueryExpander.builder()
                .chatClientBuilder(builder)
                .includeOriginal(false)
                .numberOfQueries(3)
                .build();
        List<Query> queries = expander.expand(new Query(message));
        queries.stream().forEach(item -> System.out.println(item.text()));
        //调用大模型
        List<Message> list = queries.stream().map(query -> new UserMessage(query.text())).collect(Collectors.toList());
        String content = chatClient.prompt(new Prompt(list))
                .call().content();
        return content;
    }

    @GetMapping("/rag/retrieval-retriever")
    @Operation(summary = "rag检索处理-文档检索")
    public List<Document>  rag5(@RequestParam(defaultValue = "根据已经提供的资料，描述下Spring AI是什么？") String message) {
        //构建示例文档
        List<Document> documents = new ArrayList<>();
        documents.add(new Document( "Spring AI 是一个基于 Spring 生态系统的开源项目，旨在简化人工智能（AI）和机器学习（ML）技术在 Java 应用中的集成。它提供了与各种 AI/ML 框架、工具和服务的无缝集成，使开发者能够轻松地在 Spring 应用中引入 AI 功能。" +
                "主要功能" +
                "1. 模型集成：支持与主流 AI 模型（如 OpenAI、Hugging Face 等）的集成。" +
                "2. 数据处理：提供数据预处理和特征工程工具，便于准备训练数据。" +
                "3. 训练与推理：支持模型训练和推理，开发者可直接在应用中调用训练好的模型。" +
                "4. 自动化：包含 AutoML 功能，自动选择最佳模型和超参数。" +
                "5. 可扩展性：设计灵活，易于扩展和自定义。" +
                "使用场景" +
                "1.自然语言处理（NLP）：如文本分类、情感分析。" +
                "2.计算机视觉：如图像识别、目标检测。" +
                "3.推荐系统：如个性化推荐。" +
                "4.预测分析：如时间序列预测。" +
                "总结" +
                "Spring AI 为 Java 开发者提供了便捷的 AI/ML 集成工具，帮助快速构建智能应用。",
                Map.of("year", 2025,"month",2)));
        //把文档添加到向量数据库
        vectorStore.add(documents);
        //构建复杂的filterExpress
        FilterExpressionBuilder bd = new FilterExpressionBuilder();
        Filter.Expression fliterExpress = bd.and(
                bd.eq("year", 2025),
                bd.in("content", "Spring AI")
        ).build();
        //配置文档检索器
        VectorStoreDocumentRetriever retriever = VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .similarityThreshold(0.5)
                .topK(3)
                //经测试，在vectorstore中设置了元数据的字段，这个地方筛选才会生效;否则使用filterExpression会检索不到数据
                //.filterExpression(new FilterExpressionBuilder().eq("year",2025).build())
                .filterExpression(fliterExpress)
                .build();
        Query query=Query.builder()
                .text(message)
                .build();
        List<Document> list = retriever.retrieve(query);
        return list;
        //创建上下文感知的查询增强器
//        RetrievalAugmentationAdvisor advisor = RetrievalAugmentationAdvisor.builder()
//                .queryAugmenter(ContextualQueryAugmenter.builder()
//                        .allowEmptyContext(true)// 允许上下文为空，避免NPE
//                        .build())
//                .documentRetriever(retriever)
//                .build();
//        //调用大模型
//        String content = chatClient.prompt()
//                .user(message)
//                .advisors(advisor)
//                .call()
//                .content();
//        return content;
    }
    @GetMapping("/rag/retrieval-joiner")
    @Operation(summary = "rag检索处理-文档合并")
    public List<Document> rag6(@RequestParam(defaultValue = "Spring AI 支持哪些功能") String message) {
        Map<Query, List<List<Document>>> documentsForQuery = new HashMap<>();
        List<Document> documents = new ArrayList<>();
        // 测试的document数据，实际中应该是从多个数据源拿到的数据
//        documents.add(new Document( "Spring AI 是一个强大的框架",  Map.of("year", "2025","month","02")));
//        documents.add(new Document("它支持RAG和生成模型",Map.of("year","2025","month","02")));
//        documents.add(new Document("它支持文生图、文生音频、文生视频等",Map.of("year","2025","month","03")));
        Document d1 = Document.builder()
                .id("a")
                .text("Spring AI支持RAG和生成模型")
                .metadata(Map.of("year", "2025", "month", "02"))
                .score(0.8).build();
        Document d2 =Document.builder()
                .id("b")
                .text("Spring AI是一个强大的框架")
                .metadata(Map.of("year", "2025", "month", "02"))
                .score(0.7).build();
        Document d22 =Document.builder()
                .id("b")
                .text("Spring AI是一个强大的框架")
                .metadata(Map.of("year", "2025", "month", "02"))
                .score(0.7).build();
        documents.add(d1);
        documents.add(d2);
        documents.add(d22);
        documentsForQuery.put(new Query(message), List.of(documents));
        DocumentJoiner documentJoiner = new ConcatenationDocumentJoiner();
        List<Document> list = documentJoiner.join(documentsForQuery);
        return list;
//        vectorStore.add(list);
//        //配置文档检索器
//        VectorStoreDocumentRetriever retriever = VectorStoreDocumentRetriever.builder()
//                .vectorStore(vectorStore)
//                .similarityThreshold(0.5)
//                .topK(3)
//                .filterExpression(()->new FilterExpressionBuilder()
//                        .eq("year", "2025")
//                        .build())
//                .build();
//        //创建上下文感知的查询增强器
//        RetrievalAugmentationAdvisor advisor = RetrievalAugmentationAdvisor.builder()
//                .queryAugmenter(ContextualQueryAugmenter.builder()
//                        .allowEmptyContext(true)// 允许上下文为空，避免NPE
//                        .build())
//                .documentRetriever(retriever)
//                .build();
//        //调用大模型
//        String content = chatClient.prompt()
//                .user(message)
//                .advisors(advisor)
//                .call()
//                .content();
//        return content;
    }
    @GetMapping("/rag/post-retrieval-ranker")
    @Operation(summary = "rag检索处理-文档排序")
    public List<Document> rag7(@RequestParam(defaultValue = "Spring AI 支持哪些功能") String message) {
        List<Document> documents = new ArrayList<>();
        // 测试的document数据，实际中应该是从多个数据源拿到的数据
        documents.add(new Document("Spring AI 是一个强大的框架",  Map.of("year", "2025","month","02")));
        documents.add(new Document("它支持RAG和生成模型",Map.of("year","2025","month","02")));
        documents.add(new Document("RAG 是一种结合检索和生成的技术",Map.of("year","2025","month","03")));

        DocumentRanker documentRanker=new DocumentRanker() {
            @Override
            public List<Document> rank(Query query, List<Document> documents) {
                //自定义ranker的逻辑，这里用文档内容的长度进行排序
                documents.sort(Comparator.comparingInt(o -> o.getText().length()));
                return documents;
            }
        };
        List<Document> list = documentRanker.rank(new Query(message), documents);
        return list;
//        vectorStore.add(list);
        //配置文档检索器
//        VectorStoreDocumentRetriever retriever = VectorStoreDocumentRetriever.builder()
//                .vectorStore(vectorStore)
//                .similarityThreshold(0.5)
//                .topK(3)
//                .filterExpression(()->new FilterExpressionBuilder()
//                        .eq("year", "2025")
//                        .build())
//                .build();
//        //创建上下文感知的查询增强器
//        RetrievalAugmentationAdvisor advisor = RetrievalAugmentationAdvisor.builder()
//                .queryAugmenter(ContextualQueryAugmenter.builder()
//                        .allowEmptyContext(true)// 允许上下文为空，避免NPE
//                        .build())
//                .documentRetriever(retriever)
//                .build();
//        //调用大模型
//        String content = chatClient.prompt()
//                .user(message)
//                .advisors(advisor)
//                .call()
//                .content();
//        return content;
    }
    @GetMapping("/rag/post-retrieval-selector")
    @Operation(summary = "rag检索处理-文档选择")
    public List<Document>  rag8(@RequestParam(defaultValue = "Spring AI 支持哪些功能") String message) {
        List<Document> documents = new ArrayList<>();
        // 测试的document数据，实际中应该是从多个数据源拿到的数据
        documents.add(new Document("Spring AI 是一个强大的框架",  Map.of("year", "2025","month","02")));
        documents.add(new Document("它支持RAG和生成模型",Map.of("year","2025","month","02")));
        documents.add(new Document("RAG 是一种结合检索和生成的技术",Map.of("year","2025","month","03")));

        DocumentSelector documentSelector=new DocumentSelector() {
            @Override
            public List<Document> select(Query query, List<Document> documents) {
                //自定义selector的逻辑，不包含RAG内容的提出掉
                documents.removeIf(document -> !document.getText().contains("RAG"));
                return documents;
            }
        };
        List<Document> list = documentSelector.select(new Query(message), documents);
        return  list;
//        vectorStore.add(list);
//        //配置文档检索器
//        VectorStoreDocumentRetriever retriever = VectorStoreDocumentRetriever.builder()
//                .vectorStore(vectorStore)
//                .similarityThreshold(0.5)
//                .topK(3)
//                .filterExpression(()->new FilterExpressionBuilder()
//                        .eq("year", "2025")
//                        .build())
//                .build();
//        //创建上下文感知的查询增强器
//        RetrievalAugmentationAdvisor advisor = RetrievalAugmentationAdvisor.builder()
//                .queryAugmenter(ContextualQueryAugmenter.builder()
//                        .allowEmptyContext(true)// 允许上下文为空，避免NPE
//                        .build())
//                .documentRetriever(retriever)
//                .build();
//        //调用大模型
//        String content = chatClient.prompt()
//                .user(message)
//                .advisors(advisor)
//                .call()
//                .content();
//        return content;
    }
    @GetMapping("/rag/post-retrieval-compression")
    @Operation(summary = "rag检索处理-文档压缩")
    public List<Document>  rag9(@RequestParam(defaultValue = "Spring AI 支持哪些功能") String message) {
        List<Document> documents = new ArrayList<>();
        // 测试的document数据，实际中应该是从多个数据源拿到的数据
        documents.add(new Document("Spring AI 是一个强大的框架",  Map.of("year", "2025","month","02")));
        documents.add(new Document("它支持RAG和生成模型",Map.of("year","2025","month","02")));
        documents.add(new Document("RAG 是一种结合检索和生成的技术",Map.of("year","2025","month","03")));

        DocumentCompressor documentCompressor=new DocumentCompressor() {
            @Override
            public List<Document> compress(Query query, List<Document> documents) {
                //自定义compressor的逻辑，截取内容的前6个字符
                List<Document> collect = documents.stream().map(item -> {
                    String substring = item.getText().substring(0, 6);
                    return new Document(substring, item.getMetadata());
                }).collect(Collectors.toList());
                return collect;
            }
        };
        List<Document> list = documentCompressor.compress(new Query(message), documents);
        return  list;
//        vectorStore.add(list);
//        //配置文档检索器
//        VectorStoreDocumentRetriever retriever = VectorStoreDocumentRetriever.builder()
//                .vectorStore(vectorStore)
//                .similarityThreshold(0.5)
//                .topK(3)
//                .filterExpression(()->new FilterExpressionBuilder()
//                        .eq("year", "2025")
//                        .build())
//                .build();
//        //创建上下文感知的查询增强器
//        RetrievalAugmentationAdvisor advisor = RetrievalAugmentationAdvisor.builder()
//                .queryAugmenter(ContextualQueryAugmenter.builder()
//                        .allowEmptyContext(true)// 允许上下文为空，避免NPE
//                        .build())
//                .documentRetriever(retriever)
//                .build();
//        //调用大模型
//        String content = chatClient.prompt()
//                .user(message)
//                .advisors(advisor)
//                .call()
//                .content();
//        return content;
    }
    @GetMapping("/rag/advisor")
    @Operation(summary = "rag-检索增强顾问")
    public String  rag10(@RequestParam(defaultValue = "根据已经提供的资料，描述下Spring AI主要功能？") String message) {
        //创建上下文感知的查询增强器
        RetrievalAugmentationAdvisor advisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .vectorStore(vectorStore)
                        .similarityThreshold(0.5)
                        .topK(3)
                        .filterExpression(new FilterExpressionBuilder().eq("year",2025).build())
                        .build())//检索文档
                .queryTransformers(RewriteQueryTransformer.builder()
                        .chatClientBuilder(builder.build().mutate())
                        .build())
//                .queryAugmenter(ContextualQueryAugmenter.builder()
//                        .allowEmptyContext(true)// 允许上下文为空，避免NPE
//                        .build())
                .build();
         //调用大模型
        String content = chatClient.prompt()
                .advisors(advisor)
                .user(message)
                .call()
                .content();
        return content;
    }
}
